vllm-ascend PR #18671 解读：GLM V1 SP 布局贯穿 DSA-CP——十行 attention 到昇腾实现全链

1. PR 档案

2. 问题：布局来回倒腾（改动前）

先理清三方角色，关键是看懂 SP 怎么切。TP=16 的组里，3504 个 token（3500 的 prompt 对齐到 16 的倍数，多出的 4 个是 padding）按连续段横向切开——每卡分到 219 个，这就是模型层（DecoderLayer）手里的本地 token 分片。MoE / MLP 天然可切：各卡算各自那 219 个 token 的分片就行，谁也不需要别人的。DSA-CP attention 同样按 token 切卡：每卡只对自己那 219 个 token 的 Q 做注意力，缺的 KV 跨卡定点去取（③）。按理说模型层和 attention 拿的都是分片、布局本该一致，但原实现里对不上——老约定认为「attention 必须在全量 token 上算」，于是有了下面的 ①②④（这张切法的全貌见本节末的图 3）。

图 2 左侧是业务主线（与并行无关的算子链），右侧是并行框架挂上去的通信步骤，本 PR 只动右侧。原路径的问题：框架在 LN→注意力之间先 ① AllGather 聚成全局（[1,219,8192] → [1,3504,8192]）、attention 入口 _prepare_native_hidden_states 又 ② 切回本地段，算完 O 投影输出全局布局、框架再 ④ ReduceScatter 收回——①②④ 全是无用功，每层都来一遍；③ 取命中 KV 是 DSA-CP 本职通信。

①~⑥ 速查（后文正文与代码注释均用此编号）

3. 核心思路：本地分片贯穿三个边界

PR 立了一条布局契约。DSA-CP attention 后端声明 supports_sequence_parallel_input = True，含义是"我收本地分片，还回本地分片"。对照图 4（业务骨架与图 2 完全一致）：①② 两处冗余直接省掉；② 的防错位职责由新增的 ⑤ 输入校验 + padding 清零（纯本地操作）接替；④ 的归约职责并入 attention 内部的 ⑥ all-to-all + 单次 reduce-scatter；③ 本职通信原样保留。这样三个边界（DecoderLayer ↔ MLA 包装层 ↔ DSA-CP attention）看到的是同一份连续 TP token 分片。

3.1 为什么 ①② 能省

因为这是一对互逆操作。① AllGather 的存在理由，是上游 SP 的老约定——"attention 必须在全量 token 上算"，所以每层都要把每卡 [1,219,8192] 聚成全局 [1,3504,8192]。但 DSA-CP 的注意力天生按 token 分片算：每张卡只算自己那 219 个 token 的 Q 行，从来不需要别的卡上的 token；它唯一要的跨卡数据是命中的 KV 块，而那部分它自己通过 ③ 去拉。于是 ① 聚出来的全局 token，唯一的消费者就是紧跟着的 ②（再切回本地段）——聚了又切，切完就扔，两步互相抵消，每层每批次白付一次 AllGather 通信加两次布局整理。

3.2 为什么 ④ 能省

因为它的活儿已经有人干了。跨卡归约（把各卡的部分和加起来）这件事必须做、且只能做一次：改动后由 ⑥ 在 attention 内部完成——all_to_all 换完 head 后，数据正好是"每卡一段 token 的部分和"，紧接着一次 reduce_scatter 同时完成求和和按卡散射，一步到位。模型层原来那个 ④ 如果保留，要么重复归约（结果就错了），要么得先把数据倒腾回旧布局再散一次（通信翻倍）。所以它不是简单删掉，而是职责移进了唯一正确的位置。至于正确性：⑤ 的边界校验保证布局不可能悄悄错（错就当场抛错），FP16 的 scale 放在归约前乘，与上游语义逐位对齐——4.10 节的单测验证了新旧路径结果一致。

4. 从十行 attention 样例到昇腾实现全链

先对齐地基。第一张图先看全貌——模型架构 → 框架流程（图 5）：模型的每一件在 vllm-ascend 里住哪、谁调它。本章沿一条固定路线走，小节对应六步：4.1 架构全貌（图 5）→ 4.2 目录地图（图 6）→ 4.3 十行样例怎么被完成（流水线图 7）→ 4.4 全链时序（图 8）→ 4.5 为什么这么实现·逐步收益（图 9）→ 4.6 问题 → 4.7 这次 PR（图 10~14），4.8~4.10 三张表收拢，每张图配一段讲解。先说结论：attention 的数学一行没动；"那一堆"其实是两样东西——这份代码在框架里的原样出现，加上为了在 16 张卡上跑起来、在它周围加的通信与选址。

4.1 先看全貌：模型架构 → 框架流程

4.2 先认目录：与 attention 相关的目录，一个一个认

四个目录，一条链。worker/ 管执行与喂料，models/ 管模型结构，ops/ 管算子门面，attention/ 管注意力实现——目录名就是分工，谁也不越界。

其中两个目录要特别留意：models/ 是十行样例的调用者——层对象把 input LN、self_attn、post LN、MoE 串成一次 forward；ops/ 和 attention/ 是十行样例的实现者——门卫放行后，九步流水线把每一行数学落地。为什么实现要拆成两个模块？ops/ 管"怎么接进 vLLM 的层体系"，attention/ 管"注意力策略本身"——接法和算法是两件事，所以分家。

4.3 十行样例，vllm-ascend 原本是怎么完成的

先摆出十行样例代码——矩阵乘、RoPE、KV Cache、softmax，十来行（GQA + KV Cache 的典型写法，batch 维省略）：

# attention 朴素样例（推理版，等价于典型的 GQA+KVCache 教程实现）
def forward(x, is_prefill):                            # x: [1, seq, hidden]
    q = x @ wq                                         # L1  Q 投影
    k = x @ wk;  v = x @ wv                            # L2  K/V 投影
    q = rope(q, pos);  k = rope(k, pos)                # L3  RoPE，只转 Q/K
    k_cache = k if is_prefill else cat(k_cache, k)        # L4  KV Cache：整条历史拼在本卡
    v_cache = v if is_prefill else cat(v_cache, v)
    k_full = k_cache.repeat_interleave(4, dim=1)       # L5  GQA：2 组 KV 复用给 8 个 Q 头
    scores = q @ k_full.transpose(-1,-2) / sqrt(dk)    # L6  打分——隐含假设：全量历史都在本卡
    if is_prefill: scores = causal_mask(scores)        # L7  因果遮罩
    w = softmax(scores, dim=-1)                        # L8  归一化
    out = w @ v_full                                   # L9  加权求和
    return out @ wo                                    # L10 输出投影

看懂这张图，记住三件事。

一、十行数学一步没少。左列每一行，中列都有对应——投影、RoPE、写缓存、稀疏注意力、收口，数学原样；唯一的例外是 L5（GQA repeat），MLA 系用 latent KV 全头共享达到同样目的，所以没有独立步骤。

二、实现分散在三个模块，但只有一条调用链。models/ 发起 → ops/ 放行 → attention/ 干活；右列就是每一站的地盘和门牌（文件:行号）。

三、此时还没有 PR 的事。这就是 vllm-ascend 平时完成这份 attention 的完整样子。它什么时候出问题、这次 PR 改了它哪里，往下看。

4.4 全链时序：谁调用谁，每一棒先自己干再交棒

图 7 看的是"数学落在哪"。这份工作是谁发起、又一层层交给谁的？下面把整条调用链按时间拉开——六个类六条泳道，每个函数先标"自己干了啥"，再标"委托给谁"：

读这张图抓三个节奏。

一、委托链只有一条，没有岔路。引擎组批 → Worker 接批 → Runner 喂料 → 层外壳发起 → 门卫放行 → 后端开算；每一棒只把活交给下一棒，谁也不越级。

二、"自己干"的都是壳上的杂活，"委托"出去的才是正餐。Worker 自己只做监控与流水线收发；Runner 自己只做 pad 和组元数据；把十行样例真正算出来的，是最右边后端的那串自转圈——九步流水线。

三、门卫是一趟直通车。torch.ops 这层 custom op 让编译器到此止步：层外壳不感知硬件，后端不感知模型，两边都只认门卫；后端是谁，加载期就由 resolve_sfa_impl 四选一定死了。

4.5 为什么这么实现：逐步收益

图 7 回答了"落在哪"，图 8 回答了"谁调用谁"，还剩最后一个问题：每一步为什么这么实现？各自省的是什么？下面把图 7 中列的流水线原样搬到最左列，向右逐层展开——怎么做（带门牌）→ 为什么这么设计 → 收益点：

逐行看收益，"省"分三种，省的东西不一样。

省显存的：低秩 latent 让 KV Cache 不随头数翻倍（步3/无此步两行）；slot 页表定位省掉整块 cat 的大拷贝（步4/步7 一行）。

省带宽的：融合算子让中间结果不落显存——投影+rmsnorm+rope+写缓存一个算子（步4），QKᵀ→softmax→·V 一个算子（步8 三行）；分数矩阵从头到尾不写回显存。

省算力的：稀疏注意力只算 indexer 选中的 top-2048 块，不扫全量历史；序列越长，省得越多。

也有一笔反向账要认：DSA-CP 每卡都存全量 KV（16 份）——这是 ③ 跨卡取命中块的前提，拿显存换正确性；这笔账 PR 之后依然在（PR 不动 ③）。

PR 在这张图里的位置：①④ 在层外壳（第一行），⑤ 在输入准备，⑥ 在收口（最后一行）——三处全在"壳"上，九步数学一步没动。往下就看 PR 怎么拆壳。

4.6 开了 SP + DSA-CP，原本的跑法出了问题

L6 之所以写得那么短，是因为 L4 已经把全量历史 KV cat 进了同一张卡——"所有 token 的 KV 我都拿得到"是单卡默认成立的隐含假设。16 卡把 token 切成 16 段后，每张卡只分到 219 个 token，这个假设碎了。上游 vLLM 的解法很直接：每层先把 token 聚回全量，让 attention 照旧在全量上跑，算完再散回去——这就是第 2 章的 ① 和 ④（上游 deepseek_v2.py:1315，有删节）：

# 上游每层都在十行样例外面套了这两行（有删节）
if input_is_sequence_parallel:
    hidden = tensor_model_parallel_all_gather(hidden, 0)   # ① 219 → 3504：把假设"买"回来，每层白付一次
hidden = self.self_attn(positions, hidden, ...)            # ← L1~L10 在这里面，照旧全量跑
...
hidden = tensor_model_parallel_reduce_scatter(hidden, 0)   # ④ 3504 → 219：再把假设"退"回去

更亏的是 L4：DSA-CP 的注意力本来就要跨卡取命中的 KV 块（③），于是 ① 聚出来的全量 token 转头又被切回本地段（②）——聚了又切，切完就扔，每层每批都来一遍。这就是第 2 章图 2 的灰块；问题不在十行数学，在"喂给它的数据布局"来回倒腾。

4.7 这次 PR：拆掉聚散，分片直通

PR 的模型层改成：分片直接进 attention，十行样例在分片上原样跑：

# PR 新增：vllm_ascend/models/deepseek_mtp.py
class AscendDSACPDecoderLayer(DeepseekV2DecoderLayer):
    def forward(self, positions, hidden_states, residual, llama_4_scaling=None):
        input_is_sequence_parallel = (self.use_sequence_parallel_moe
                                      and residual is not None
                                      and hidden_states.shape[0] != positions.shape[0])
        # 回退三条件：不是分片 / 带缩放的特殊批 / attention 契约不支持 → 原样走 4.6 的上游老路
        if (not input_is_sequence_parallel
                or llama_4_scaling is not None
                or not self.self_attn.mla_attn.supports_sequence_parallel_input):
            return super().forward(positions, hidden_states, residual, llama_4_scaling)

        hidden_states, residual = self.input_layernorm(hidden_states, residual)
        output_scale = 1.0 / self.routed_scaling_factor if hidden_states.dtype == torch.float16 else 1.0
        hidden_states = self.self_attn.mla_attn.forward_sequence_parallel(
            positions, hidden_states, output_scale=output_scale)   # 分片直进，无① 无④
        if hidden_states.dtype == torch.float16 and self.layer_idx == 0:
            residual *= 1.0 / self.routed_scaling_factor           # FP16 防溢出：第 0 层 residual 同步缩放
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
        hidden_states = self.mlp(hidden_states, already_sequence_parallel=True)
        return hidden_states, residual

这张图就是整个 PR。左列是原本（开了 SP+DSA-CP 的跑法），中列是 PR 的做法，逐行对着看。

改的是三处。层外壳：不再 ① 聚全量 / ④ 散回，分片直进（新类 AscendDSACPDecoderLayer，住 models/）；输入侧：新增 ⑤ 只查不切（四项断言，防静默错位）；输出侧：新增 ⑥ 两路收口（路 A 全权重零通信，路 B all_to_all + 单次 reduce_scatter——归约只做一次，FP16 缩放在归约前乘）。

不动的是三处。③ 跨卡取命中 KV、步 8 的稀疏注意力数学、门卫与闸门的兜底——任何一级条件不满足，整条链回左列老路，行为与改前完全一致。

4.8 对照总表：十行每一行，在这次 PR 里去哪了

前面把"原本怎么跑"和"PR 怎么改"都走完了，这张总表把十行样例与 PR 新路径的对应关系收拢到一起（文件与行号全部经本地源码实测）：

4.9 PR 落点总表

4.10 测试（5 个文件，覆盖面很扎实）

📍 链上定位 · 横跨第 4 章全部改动点 · 角色：给整条新链路织回归网（正确性、回退、缩放语义逐项钉死）

5. 前后代码对照（带注释）

"改后"代码均摘自 PR 实际 diff（有删节，只保留主干）；"改前"的模型层为上游 vLLM 真实代码（DeepseekV2DecoderLayer.forward，deepseek_v2.py，有删节）。绿色加粗行为本 PR 新增逻辑。注释里的 ①~⑥ 序号对应图 2 / 图 4 的通信步骤标记（① AllGather、② 再切本地段、③ 取命中 KV、④ ReduceScatter、⑤ 输入校验、⑥ all-to-all+单次RS）。

5.1 模型层 forward：从"AG → 算 → RS"到布局直通

# 上游约定：attention 要"全量 token"，SP 的 MoE 要"本地分片"——两头迁就，每层倒腾一次
def forward(self, positions, hidden_states, residual, llama_4_scaling=None):
    full_num_tokens = positions.shape[0]                       # 全量 token 数（3504）
    input_is_sequence_parallel = (self.use_sequence_parallel_moe
                                  and residual is not None
                                  and hidden_states.shape[0] != full_num_tokens)  # 219 ≠ 3504
    hidden_states, residual = self.input_layernorm(hidden_states, residual)
    if input_is_sequence_parallel:
        hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)  # ← 图2-① 分片聚成全局，每层白做
        hidden_states = hidden_states[:full_num_tokens]
    hidden_states = self.self_attn(positions, hidden_states, llama_4_scaling)  # attention 在全量上算
    ...
    hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0)      # ← 图2-④ 全局散回分片，每层白做
    hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
    hidden_states = self.mlp(hidden_states, already_sequence_parallel=True)
    return hidden_states, residual

class AscendDSACPDecoderLayer(DeepseekV2DecoderLayer):
    def forward(self, positions, hidden_states, residual, llama_4_scaling=None):
        # 判定输入是否已是"本地 token 分片"：SP MoE 开启 + residual 存在
        # + 本地 token 数 != 全局位置数（分片的典型特征）
        input_is_sequence_parallel = (
            self.use_sequence_parallel_moe and residual is not None
            and hidden_states.shape[0] != positions.shape[0]
        )
        # 首层 replicated 输入 / 不支持的投影契约 → 原样走上游，行为不变
        if (not input_is_sequence_parallel
                or llama_4_scaling is not None
                or not self.self_attn.mla_attn.supports_sequence_parallel_input):
            return super().forward(positions, hidden_states, residual, llama_4_scaling)

        # 直通链（图2）：省掉 ①④，全程无 AllGather / ReduceScatter
        hidden_states, residual = self.input_layernorm(hidden_states, residual)
        # FP16 防溢出：把 routed_scaling 缩放交给 attention 在 TP 归约"前"做
        output_scale = 1.0 / self.routed_scaling_factor if hidden_states.dtype == torch.float16 else 1.0
        # 收分片、还分片（走 4.7 的新入口）：①④ 在此省掉，归约见 5.3 的 ⑥
        hidden_states = self.self_attn.mla_attn.forward_sequence_parallel(
            positions, hidden_states, output_scale=output_scale)
        # FP16 且第 0 层：residual 同步缩放，逐位对齐上游 overflow 保护
        if hidden_states.dtype == torch.float16 and self.layer_idx == 0:
            residual *= 1.0 / self.routed_scaling_factor
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
        # 告诉 MLP "输入已是 SP 分片"，MLP 不再做布局整理
        hidden_states = self.mlp(hidden_states, already_sequence_parallel=True)
        return hidden_states, residual

5.2 输入准备：从"全局切段"到"校验分片"

# 进来的是"全局 token"（上游已 AllGather），在这里切出本 rank 的一段
# —— 图2-②：刚聚的全局，转头又切回本地，正是要消掉的浪费
return hidden_states[context.local_start : context.local_end_with_pad]

def _prepare_sequence_parallel_hidden_states(self, hidden_states, attn_metadata):
    context = attn_metadata.dsa_cp_context
    tp_group = get_tp_group()
    local_tokens = hidden_states.shape[0]
    expected_start = tp_group.rank_in_group * local_tokens
    # ⑤ 新增（接替图2-② 的位置）：校验"本 rank 那段"，不一致立刻抛错（防静默错位）
    if (context.num_tokens != attn_metadata.num_input_tokens
            or context.num_tokens_pad != local_tokens * self.tp_size
            or context.local_start != expected_start
            or context.local_end_with_pad != expected_start + local_tokens):
        raise RuntimeError("SFA DSA-CP input does not match its sequence-parallel shard: ...")
    # Graph/DP 的 padding 区清零；先 clone 再清，避免污染 residual 的共享存储
    valid_tokens = max(0, min(local_tokens, context.num_tokens - context.local_start))
    if valid_tokens < local_tokens:
        hidden_states = hidden_states.clone()
        hidden_states[valid_tokens:].zero_()
    return hidden_states

5.3 输出收口 o_proj：路 A 全权重直算，路 B all-to-all + 单次 RS

def _finalize_sequence_parallel_o_proj(self, attn_output, output, gather_full_o_proj,
                                       attn_metadata, output_scale):
    if gather_full_o_proj:
        # 路 A：full head + full weight —— 直接算出完整本地 token，零集合通信
        with self._use_full_o_proj_weights():
            output.copy_(self._apply_o_proj_full_weight(attn_output))
        if output_scale != 1.0:
            output.mul_(output_scale)
    else:
        # 路 B（⑥，常规 TP head 分片）：attn_output 是 [本地token, 全head]，
        # view 成 [本地token, tp, head/tp] 后 permute，使 tp 维落到最前；
        # all_to_all 后变成 [全token, 本地head] —— 每张卡拿到全部 token 的自己的 head
        send = (attn_output.view(-1, self.tp_size, self.num_heads * self.v_head_dim)
                .permute(1, 0, 2).reshape(-1, self.num_heads * self.v_head_dim))
        sharded_output = torch.empty_like(send)
        tp_group = get_tp_group()
        torch.distributed.all_to_all_single(sharded_output, send, group=tp_group.device_group)
        projected_output = self.o_proj(sharded_output)[0]      # 行并行：结果未归约
        if output_scale != 1.0:
            projected_output.mul_(output_scale)                # FP16：归约"前"缩放
        # ⑥ 单次 reduce_scatter = 归约 + 散射 → 本地分片（图2-④ 在此消失）
        torch.distributed.reduce_scatter_tensor(output, projected_output.contiguous(),
                                                group=tp_group.device_group)
    context = attn_metadata.dsa_cp_context
    valid_tokens = max(0, min(output.shape[0], context.num_tokens - context.local_start))
    if valid_tokens < output.shape[0]:
        output[valid_tokens:].zero_()                          # 尾部 padding 清零
    return output

5.4 能力声明与探测：安全阀

# sfa_cp.py —— 两个后端一真一假，DCP 明确不硬吃契约
class AscendSFADSACPImpl(OProjWeightSwitchMixin, AscendSFAImpl):
    supports_sequence_parallel_input = True   # DSA-CP：TP 组内按 token 分片 → 支持

class AscendSFADCPImpl(DCPImplMixin, AscendSFAImpl):
    supports_sequence_parallel_input = False  # DCP：query/head 两轴分别恢复 → 不支持

# ops/mla.py —— MLA 包装层的四条件探测，任一不满足自动回退原路径
@property
def supports_sequence_parallel_input(self) -> bool:
    impl = self.mla_attn.impl
    return (
        self.output_token_shard_size == 1                          # 输出按 token 分片
        and getattr(impl, "supports_sequence_parallel_input", False)  # 后端声明支持
        and impl.o_proj.custom_op is None                          # 无自定义通信算子
        and not impl.o_proj.reduce_results                         # o_proj 不自带 reduce
    )

5.5 模型选择：四个条件全中才换新路径

class AscendGlmMoeDsaForCausalLM(GlmMoeDsaForCausalLM):
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        if (not vllm_config.use_v2_model_runner                        # 条件1：V1 runner（use_v2=False）
                and vllm_config.parallel_config.pipeline_parallel_size == 1  # 条件2：PP=1
                and vllm_config.parallel_config.use_sequence_parallel_moe    # 条件3：SP MoE 开启
                and enable_dsa_cp()):                                  # 条件4：enable_dsa_cp
            self.model_cls = AscendDSACPModel    # 全中 → 换布局直通的原生模型
        super().__init__(vllm_config=vllm_config, prefix=prefix)
        # 不满足任一条件 → 走原 DeepseekV2Model，行为完全不变

6. 与我们排查问题的关系（重点）

它优化的正是 SP + DSA-CP 路径上的冗余通信，作用点在 prefill/decode 混合批次的每一层。我们之前测得：enable_prefill_mc2 无收益（TTFT 739→753-768 持平），1009 相对 0815 的 TTFT 差距锁定在 DispatchFFNCombine kernel 回退——而这个 PR 是另一条独立的优化轴（省模型层与 attention 边界的 AG/RS），两者不冲突，理论上可叠加。

PR 报告的单次候选测量：TTFT 均值 529.3ms、TPOT 21.4ms、输出吞吐 348.7 tokens/s、零 KV 抢占；GSM8K 4-shot 全量 1319 题 97.57%。

7. 收益、风险与待验证

8. 术语速查

来源：github.com/vllm-project/vllm-ascend PR #18671（draft，base: releases/v0.30.0，+923/−34，10 文件）· 本页为技术解读，数据均引自 PR 原文